package com.pg.agent.engine.client.model.llm;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.ObjectUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;

import com.pg.agent.engine.client.common.AgentConstant;
import com.pg.agent.engine.client.model.AIAgentSession;
import com.pg.agent.engine.client.model.Function;
import com.pg.agent.engine.client.model.sop.Memory;
import com.pg.base.component.PgAssert;
import com.pg.base.util.PgHttpUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * llm 工具类
 *
 * @author 张俊敏
 * @date 2023/11/15 16:42
 */
@Slf4j
public class ComponentLLMUtil {


    /**
     * 获取大模型，如果传入大模型为空，则取默认大模型
     *
     * @param llm
     * @return
     */
    public static ComponentLLM getLLMOrDefaultLLM(ComponentLLM llm) {

        if (llm == null) {
            llm = new ComponentLLM();
            llm.setUrl("https://itest.clife.net/module/gpt-4/getChatCompletions");
            llm.setRequestMethod("POST");
        }
        return llm;
    }

    public static LLMCallInfo call(LLMCallInfo lLMCallInfo, ComponentLLM llm, AIAgentSession session, List<Memory> chatHistory, String systemPrompt, String lastPrompt) {

        LLMCallInfo callInfo = getResponse(lLMCallInfo, llm, chatHistory, systemPrompt, lastPrompt);

        // 输出日志
        logInfo(lLMCallInfo);

        session.getInfoList().add(callInfo);

        return callInfo;
    }

    public static void logInfo(LLMCallInfo lLMCallInfo) {

        log.info("\n{}({})********大模型提示词组装如下*******\n{}\n{},\n********大模型结果如下*******\n{} \n***************************结束****************************",
                lLMCallInfo.getCallType(), lLMCallInfo.getCallName(), lLMCallInfo.getSystemPrompt(), lLMCallInfo.getUserPrompt(), lLMCallInfo.getResponse());

    }

    /**
     * 获取LLM请求响应数据
     *
     * @param chatHistory  历史记忆
     * @param systemPrompt 系统场景对话
     * @param lastPrompt   输出提示
     * @return 模型返回结果
     */
    private static LLMCallInfo getResponse(LLMCallInfo lLMCallInfo, ComponentLLM llm, List<Memory> chatHistory, String systemPrompt, String lastPrompt) {

        //对话提示语：如：请尽量简明扼要地答复,控制在100个字内
        JSONArray messages = new JSONArray();
        JSONObject messagesRole = new JSONObject();
        messagesRole.put("role", "system");
        messagesRole.put("content", StringUtils.trimToEmpty(systemPrompt));
        messages.add(messagesRole);

        String systemPromptLog = "角色：system；内容：" + StringUtils.trimToEmpty(systemPrompt);
        // 获取最后10条数据
        if (CollectionUtil.isNotEmpty(chatHistory) && chatHistory.size() > AgentConstant.MAX_CHAT_HISTORY) {
            chatHistory = chatHistory.subList(chatHistory.size() - AgentConstant.MAX_CHAT_HISTORY, chatHistory.size());
        }

        StringBuilder promptSb = new StringBuilder("");
        if (chatHistory != null) {
            // 对历史对话记录进行倒序组装成发送给大模型数据格式
            for (Memory memory : chatHistory) {
                JSONObject temp = new JSONObject();
                temp.put("role", "user");
                temp.put("content", memory.getContent());
                messages.add(temp);
                promptSb.append("角色：user   内容：").append(memory.getContent());
            }
        }

        // 在最后追加提示语，返回什么结果提示
        if (ObjectUtil.isNotEmpty(lastPrompt)) {
            if (messages.size() >= 1) {
                JSONObject jsonObject = messages.getJSONObject(messages.size() - 1);
                jsonObject.put("content", jsonObject.getString("content") + lastPrompt);
                promptSb.append("。\n\n最后输出要求：").append(lastPrompt);
            }
        }

        String response = null;
        try {
            lLMCallInfo.setCallStartTime(new Date());
            long time = System.currentTimeMillis();
            JSONObject jsonObject = new JSONObject();
            jsonObject.put("messages", messages);
            String type = lLMCallInfo.getCallType();
            if ("State切换".equals(type) || "Agent切换".equals(type) || "AgentComponent".equals(type)) {
                jsonObject.put("temperature", 0);
            }
            response = chatCompletions(llm, jsonObject);

            lLMCallInfo.setCallEndTime(new Date());
            lLMCallInfo.setTimeLong(System.currentTimeMillis() - time);

        } catch (Exception e) {
            log.error("调用模型报错：", e);
        }
        //log.info("response：{}", response);
        // 判断模型，对应返回不同的数据结构
        JSONObject jsonObject = JSON.parseObject(response);
        if (jsonObject != null && ObjectUtil.equal(jsonObject.getInteger("code"), 0)) {
            JSONObject resultObj = jsonObject.getJSONArray("data").getJSONObject(0);
            String dataType = resultObj.getString("dataType");
            if ("image".equals(dataType)) {
                lLMCallInfo.setImageList(resultObj.getJSONArray("content").toJavaList(String.class));
            } else if ("text".equals(dataType)) {
                response = resultObj.getString("content");
            } else {
                response = resultObj.getString("content");
            }
            if (resultObj.getJSONObject("functionCall") != null) {
                String argument = resultObj.getJSONObject("functionCall").getString("arguments");
                lLMCallInfo.setArguments(argument);
            }
        } else {
            PgAssert.bizError(true, 1000, "大模型数据结构错误：" + response);
        }

        lLMCallInfo.setModel(llm.getModel());
        lLMCallInfo.setResponse(response);
        lLMCallInfo.setSystemPrompt(systemPromptLog);
        lLMCallInfo.setUserPrompt(promptSb.toString());

        return lLMCallInfo;
    }

    /**
     * 构建大模型对接对象和完成对话
     *
     * @param llm
     * @param jsonObject
     * @return
     */
    private static String chatCompletions(ComponentLLM llm, JSONObject jsonObject) {

        String result = null;
        Map<String, String> heard = new HashMap<>();
        heard.put("Content-Type", "application/json");
        //调用完成对话
        try {
            if (PgHttpUtil.GET.equals(llm.getRequestMethod())) {
                result = PgHttpUtil.get(llm.getUrl(), new HashMap<>(), jsonObject, 60 * 1000);
            } else {
                result = PgHttpUtil.postBody(llm.getUrl(), new HashMap<>(), JSONObject.toJSONString(jsonObject), 60 * 1000);
            }
        } catch (Exception e) {
            String errorMsg = String.format("url[%s],param[%s]", llm.getUrl(), JSON.toJSONString(jsonObject));
            PgAssert.bizError(true, 1000, "大模型接口" + errorMsg + "请求出错,原因：" + e.getMessage());
        }
        return result;
    }

    public static void callBack(LLMCallInfo lLMCallInfo, ComponentLLM llm, AIAgentSession session, List<Memory> chatHistory, String systemPrompt, String lastPrompt, List<Function> functionList) {
        JSONArray messages = new JSONArray();
        JSONObject messagesRole = new JSONObject();
        messagesRole.put("role", "system");
        messagesRole.put("content", StringUtils.trimToEmpty(systemPrompt));
        messages.add(messagesRole);

        String systemPromptLog = "角色：system；内容：" + StringUtils.trimToEmpty(systemPrompt);

        if (functionList != null) {
            systemPromptLog = systemPromptLog + "\n参数：" + JSON.toJSONString(functionList);
        }

        // 获取最后10条数据
        if (CollectionUtil.isNotEmpty(chatHistory) && chatHistory.size() > AgentConstant.MAX_CHAT_HISTORY) {
            chatHistory = chatHistory.subList(chatHistory.size() - AgentConstant.MAX_CHAT_HISTORY, chatHistory.size());
        }

        StringBuilder promptSb = new StringBuilder("");
        if (chatHistory != null) {
            // 对历史对话记录进行倒序组装成发送给大模型数据格式
            for (Memory memory : chatHistory) {
                JSONObject temp = new JSONObject();
                temp.put("role", "user");
                temp.put("content", memory.getContent());
                messages.add(temp);
                promptSb.append("角色：user，内容：").append(memory.getContent());
            }
        }

        // 在最后追加提示语，返回什么结果提示
        if (ObjectUtil.isNotEmpty(lastPrompt)) {

            if (messages.size() >= 1) {
                JSONObject jsonObject = messages.getJSONObject(messages.size() - 1);
                jsonObject.put("content", jsonObject.getString("content") + lastPrompt);
            }

        }

        String response = null;

        try {
            lLMCallInfo.setCallStartTime(new Date());
            long time = System.currentTimeMillis();
            JSONObject jsonObject = new JSONObject();
            jsonObject.put("messages", messages);
            jsonObject.put("functions", functionList);
            response = chatCompletions(llm, jsonObject);

            lLMCallInfo.setCallEndTime(new Date());
            lLMCallInfo.setTimeLong(System.currentTimeMillis() - time);

        } catch (Exception e) {
            log.error("调用模型报错：", e);
        }

        // 判断模型，对应返回不同的数据结构
        JSONObject jsonObject = JSON.parseObject(response);
        if (jsonObject != null && jsonObject.getInteger("code") == 0) {
            JSONObject data = jsonObject.getJSONArray("data").getJSONObject(0);
            response = data.getString("content");
            lLMCallInfo.setResponse(response);
            if (data.getJSONObject("functionCall") != null) {
                String argument = data.getJSONObject("functionCall").getString("arguments");
                lLMCallInfo.setArguments(argument);
                if (ObjectUtil.isNotEmpty(argument)) {
                    lLMCallInfo.setResponse(argument);
                }
            }
        } else {
            log.info("第三方接口请求出错,原因：{}", response);
        }

        lLMCallInfo.setModel(llm.getModel());
        lLMCallInfo.setSystemPrompt(systemPromptLog);
        lLMCallInfo.setUserPrompt(promptSb.toString());
        session.getInfoList().add(lLMCallInfo);

        // 输出日志
        logInfo(lLMCallInfo);

    }

    public static LLMCallInfo callConvert(LLMCallInfo llmCallInfo, ComponentLLM llm, AIAgentSession session, String userPrompt, String systemPrompt) {
        JSONArray messages = new JSONArray();

        JSONObject systemMessagesRole = new JSONObject();
        systemMessagesRole.put("role", "system");
        systemMessagesRole.put("content", systemPrompt);
        JSONObject userMessagesRole = new JSONObject();
        userMessagesRole.put("role", "user");
        userMessagesRole.put("content", userPrompt);
        messages.add(systemMessagesRole);
        messages.add(userMessagesRole);

        String response = null;

        try {
            llmCallInfo.setCallStartTime(new Date());
            long time = System.currentTimeMillis();
            JSONObject jsonObject = new JSONObject();
            jsonObject.put("messages", messages);
            jsonObject.put("temperature", 0);

            response = chatCompletions(llm, jsonObject);

            llmCallInfo.setCallEndTime(new Date());
            llmCallInfo.setTimeLong(System.currentTimeMillis() - time);

        } catch (Exception e) {
            log.error("调用模型报错：", e);
        }

        // 判断模型，对应返回不同的数据结构
        JSONObject jsonObject = JSON.parseObject(response);
        if (jsonObject != null && ObjectUtil.equal(jsonObject.getInteger("code"), 0)) {
            llmCallInfo.setResponse(response);
        } else {
            PgAssert.bizError(true, 1000, "大模型数据结构错误：" + response);
        }

        llmCallInfo.setModel(llm.getModel());
        llmCallInfo.setSystemPrompt(systemPrompt);
        llmCallInfo.setUserPrompt(userPrompt);
        session.getInfoList().add(llmCallInfo);

        // 输出日志
        logInfo(llmCallInfo);

        return llmCallInfo;
    }


    public static String extractResult(String response,List<String> resultList) {

        String type ="end";

        String targetStr = get_content_between_a_b("<" + type + ">", "</" + type + ">", response);

        if(ObjectUtil.isEmpty(targetStr)){
            for (String result : resultList) {
                if (response.contains(result)) {
                    targetStr = result;
                    break;
                }
            }
        }
        return targetStr;
    }

    public static String get_content_between_a_b(String a, String b, String text) {
        if (ObjectUtil.isEmpty(text)) {
            return "";
        }
        int start = text.indexOf(a);
        int end = text.indexOf(b, start + a.length());
        if (start == -1 || end == -1) {
            return "";
        } else {
            return text.substring(start + a.length(), end);
        }
    }
}
